"""
Dataset specification

output of datasets consist of torch tensors of image, boxes and labels
    box coordinates are of percent form
    labels are torch.LongTensor
    image is loaded using PIL.Image module

preprocessing are done with torch.Tensor's
"""
import torch.utils.data

import vortex.data.transforms as VT
from .voc import PascalVOCDataset


def get_transforms(split):
    """
    ssd specific data transform.
    """
    assert split in ['train', 'trainval', 'test', 'val']
    transforms = []
    if split in ['train', 'trainval']:
        transforms.append(VT.RandomPhotometric())
        transforms.append(VT.RandomExpand())
        transforms.append(VT.RandomScaledCrop())
        transforms.append(VT.RandomHorizontalFlip())
    
    transforms.append(VT.Resize(size=(300, 300)))
    transforms.append(VT.ToTensor())
    transforms.append(VT.Normalize())

    transforms = VT.Compose(transforms)
    return transforms


def get_data(opts, train_transform=None, val_transform=None):
    if opts.NAME == 'PascalVOCDataset':
        train_dataset = PascalVOCDataset(opts.DATA_ROOT, split='train', 
                                         transform=train_transform, 
                                         labels=opts.CLASS_NAMES, 
                                         bg0=opts.BG0)
        val_dataset = PascalVOCDataset(opts.DATA_ROOT, split='test', 
                                         transform=val_transform, 
                                         labels=opts.CLASS_NAMES,
                                         bg0=opts.BG0)
        train_loader = torch.utils.data.DataLoader(train_dataset,
                                                   batch_size=opts.BATCH_SIZE, 
                                                   shuffle=True,
                                                   collate_fn=train_dataset.collate_fn,
                                                   num_workers=opts.N_WORKERS,
                                                   pin_memory=True)
        val_loader = torch.utils.data.DataLoader(val_dataset,
                                                 batch_size=opts.BATCH_SIZE, 
                                                 shuffle=False,
                                                 collate_fn=val_dataset.collate_fn,
                                                 num_workers=opts.N_WORKERS,
                                                 pin_memory=True)
        return train_loader, val_loader

    elif opts.NAME == 'UltralyticsDataset':
        pass
    elif opts.NAME == 'YoloV3Dataset':
        pass
